package org.example;

import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

public class Tree implements INode {


    int height(TreeNode node) {
        return node == null ? -1 : node.height;
    }

    void updateHeight(TreeNode node) {
        node.height = Math.max(height(node.left), height(node.right));
    }

    TreeNode rightRotate(TreeNode node) {
        TreeNode child = node.left;
        TreeNode grandChild = child.right;

        child.right = node;
        node.left = grandChild;

        updateHeight(node);
        updateHeight(child);
        return child;
    }

    TreeNode leftRotate(TreeNode node) {
        TreeNode child = node.right;
        TreeNode grandChild = child.left;

        child.left = node;
        node.right = grandChild;

        updateHeight(node);
        updateHeight(child);

        return child;
    }

    int balanceFactor(TreeNode node) {
        return height(node.left) - height(node.right);
    }

    TreeNode rotate(TreeNode node) {
        int factor = balanceFactor(node);
        if (factor > 1) {
            if (balanceFactor(node.left) < 0) {
                node.left = leftRotate(node.left);
            }
            return rightRotate(node);
        }

        if (factor < -1) {
            if (balanceFactor(node.right) > 0) {
                node.right = rightRotate(node.right);
            }

            return leftRotate(node);
        }

        return node;
    }


    TreeNode doRemove(TreeNode node,
                      int num) {
        if (node == null) {
            return null;
        }

        if (num == node.val) {
            if (node.left != null && node.right != null) {
                TreeNode temp = node.right;
                while (temp.left != null) {
                    temp = temp.left;
                }

                node.val = temp.val;
                node.right = doRemove(node.right, temp.val);
            } else if (node.left != null) {
                return node.left;
            } else if (node.right != null) {
                return node.right;
            } else {
                return null;
            }
        } else if (num > node.val) {
            node.right = doRemove(node.right, num);
        } else {
            node.left = doRemove(node.left, num);
        }

        updateHeight(node);
        return rotate(node);
    }

    TreeNode doInsert(TreeNode node,
                      int num) {
        if (node == null) {
            return new TreeNode(num);
        }

        if (num == node.val) {
            return node;
        }

        if (num > node.val) {
            node.right = doInsert(node.right, num);
        }

        if (num < node.val) {
            node.left = doInsert(node.left, num);
        }

        updateHeight(node);
        return rotate(node);
    }

    TreeNode root;

    @Test
    void insert() {

        for (int i = 0; i < 7; i++) {
            root = doInsert(root, i);
        }

        List<Integer> res = new ArrayList<>();
        dfs(root, res);
        System.out.println(res.stream().map(String::valueOf).collect(Collectors.joining()));

        for (int i = 0; i < 3; i++) {
            root = doRemove(root, i);
        }
        res = new ArrayList<>();
        dfs(root, res);
        System.out.println(res.stream().map(String::valueOf).collect(Collectors.joining()));

        root = doRemove(root, 8);
        res = new ArrayList<>();
        dfs(root, res);
        System.out.println(res.stream().map(String::valueOf).collect(Collectors.joining()));
    }

    void dfs(TreeNode node,
             List<Integer> res) {
        if (node == null) {
            return;
        }

        dfs(node.left, res);
        res.add(node.val);
        dfs(node.right, res);
    }
}
